from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta

fake = Faker('zh_CN')

# 配置
HOSPITAL_CODE = 'H001'
HOSPITAL_NAME = '西安人民医院'
CANCERS = ['肺癌', '胃癌']
DEPARTMENTS = [('D001', '内科'), ('D002', '外科')]
PATIENT_TYPES = {
    '01': '初诊患者',
    '02': '复诊患者',
    '03': '疑似患者'
}

# 抽样概率
PROB_PAST_VISIT = {
    'no_visit': 0.1,
    '0_6_month': 0.1,
    '6_12_month': 0.1,
    '12_month': 0.1
}

def random_datetime(start, end):
    """生成指定范围的随机 datetime"""
    delta = end - start
    seconds = random.randint(0, int(delta.total_seconds()))
    return start + timedelta(seconds=seconds)

def generate_pv_record(project_id, patient_id, visit_sn):
    """生成单条患者记录"""
    cancer_name = random.choice(CANCERS)
    department_code, department_name = random.choice(DEPARTMENTS)

    visit_datetime = fake.date_time_between(start_date='-2y', end_date='now')
    discharged_datetime = visit_datetime + timedelta(days=random.randint(0, 10))
    admission_time = (discharged_datetime - visit_datetime).days if discharged_datetime else 0
    have_admission_record = 1 if admission_time > 0 else 0

    # 既往就诊随机逻辑
    past_type = random.choices(
        ['no_visit', 'outpatient_only', 'admission_only', 'outpatient_adm'],
        weights=[0.1, 0.3, 0.3, 0.3],
        k=1
    )[0]

    past_outpatient_count = random.randint(0, 5) if past_type in ['outpatient_only', 'outpatient_adm'] else 0
    past_admission_count = random.randint(0, 3) if past_type in ['admission_only', 'outpatient_adm'] else 0

    past_not_visit_sn = 1 if past_type == 'no_visit' else 0
    past_outpatient_only_have = 1 if past_type == 'outpatient_only' else 0
    past_admission_only_have = 1 if past_type == 'admission_only' else 0
    past_outpatient_have_adm_have = 1 if past_type == 'outpatient_adm' else 0

    # 0~6, 6~12, 12个月风险标记
    assess_0_6_month = 1 if random.random() < PROB_PAST_VISIT['0_6_month'] else 0
    assess_12_month = 1 if random.random() < PROB_PAST_VISIT['12_month'] else 0
    create_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    sex = random.choice(['男','女'])
    age = random.randint(18, 80)
    birthday = (datetime.now() - timedelta(days=age*365)).strftime('%Y-%m-%d')
    assess_result = random.choice([0,1])
    assess_time = random_datetime(visit_datetime, datetime.now()).strftime('%Y-%m-%d %H:%M:%S')
    high_risks = random.choice([0,1])
    patient_type_code, patient_type = random.choice(list(PATIENT_TYPES.items()))

    record = {
        'id': patient_id,
        'project_id': project_id,
        'patient_id': patient_id,
        'visit_sn': visit_sn,
        'cancer_name': cancer_name,
        'hospital_code': HOSPITAL_CODE,
        'hospital_name': HOSPITAL_NAME,
        'department_code': department_code,
        'department_name': department_name,
        'discharged_datetime': discharged_datetime.strftime('%Y-%m-%d %H:%M:%S'),
        'visit_datetime': visit_datetime.strftime('%Y-%m-%d %H:%M:%S'),
        'admission_time': admission_time,
        'have_admission_record': have_admission_record,
        'past_admission_count': past_admission_count,
        'past_outpatient_count': past_outpatient_count,
        'sex': sex,
        'age': age,
        'birthday': birthday,
        'assess_result': assess_result,
        'assess_time': assess_time,
        'high_risks': high_risks,
        'patient_type_code': patient_type_code,
        'patient_type': patient_type,
        'past_not_visit_sn': past_not_visit_sn,
        'past_outpatient_only_have': past_outpatient_only_have,
        'past_admission_only_have': past_admission_only_have,
        'past_outpatient_have_adm_have': past_outpatient_have_adm_have,
        'assess_0_6_month': assess_0_6_month,
        'assess_12_month': assess_12_month,
        'create_time': create_time
    }

    return record

# 示例生成
records = []
for i in range(1, 11):  # 生成10条数据
    records.append(generate_pv_record(project_id=1, patient_id=f'PT{i:04}', visit_sn=f'V{i:04}'))

# 打印查看
for r in records:
    print(r)
